#!/usr/bin/env python3
"""
mitmproxy脚本：简化版HTTP请求记录和curl导出
确保所有请求都被立即记录，不依赖响应完成
使用方法：
1. 实时导出：mitmdump -s export_curl_simple.py --mode upstream:http://proxy:port
2. 从文件导出：mitmdump -s export_curl_simple.py -r traffic.mitm
"""

import mitmproxy.http
from mitmproxy import ctx
import json
import os
from datetime import datetime

class SimpleCurlExporter:
    def __init__(self):
        # 输出文件路径（可根据需要修改）
        self.output_file = self.get_output_path()
        self.request_count = 0
        self.init_output_file()
    
    def get_output_path(self):
        """获取输出文件路径（当前目录）"""
        return "mitmproxy_requests.txt"
    
    def init_output_file(self):
        """初始化输出文件"""
        try:
            with open(self.output_file, "w", encoding="utf-8") as f:
                f.write("=" * 80 + "\n")
                f.write("mitmproxy HTTP请求记录\n")
                f.write(f"生成时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
                f.write(f"输出路径: {os.path.abspath(self.output_file)}\n")
                f.write("=" * 80 + "\n\n")
            print(f"✅ HTTP请求将保存到: {os.path.abspath(self.output_file)}")
            ctx.log.info(f"HTTP请求将保存到: {os.path.abspath(self.output_file)}")
        except Exception as e:
            print(f"❌ 无法创建输出文件: {e}")
            ctx.log.error(f"无法创建输出文件: {e}")
    
    def request(self, flow: mitmproxy.http.HTTPFlow) -> None:
        """处理HTTP请求（立即记录）"""
        try:
            self.request_count += 1
            self.save_request(flow)
            print(f"✅ 已记录第 {self.request_count} 个请求: {flow.request.method} {flow.request.pretty_url}")
            ctx.log.info(f"已记录第 {self.request_count} 个请求: {flow.request.method} {flow.request.pretty_url}")
        except Exception as e:
            print(f"❌ 记录请求失败: {e}")
            ctx.log.error(f"记录请求失败: {e}")
    
    def save_request(self, flow: mitmproxy.http.HTTPFlow):
        """保存请求信息和curl命令"""
        try:
            with open(self.output_file, "a", encoding="utf-8") as f:
                f.write("=" * 80 + "\n")
                f.write(f"请求 #{self.request_count}\n")
                f.write(f"时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
                f.write(f"方法: {flow.request.method}\n")
                f.write(f"URL: {flow.request.pretty_url}\n")
                f.write("=" * 80 + "\n\n")
                
                # 请求头信息
                f.write("📋 请求头:\n")
                for name, value in flow.request.headers.items():
                    f.write(f"  {name}: {value}\n")
                f.write("\n")
                
                # 请求体信息（如果有）
                if flow.request.content:
                    f.write("📦 请求体:\n")
                    content_type = flow.request.headers.get("content-type", "").lower()
                    try:
                        if "application/json" in content_type:
                            json_data = json.loads(flow.request.content.decode('utf-8'))
                            json_str = json.dumps(json_data, ensure_ascii=False, indent=2)
                            f.write(f"  类型: JSON\n")
                            f.write(f"  长度: {len(flow.request.content)} bytes\n")
                            if len(json_str) > 1000:
                                f.write(f"  内容: {json_str[:1000]}...[截断]\n")
                            else:
                                f.write(f"  内容:\n{json_str}\n")
                        elif "application/x-www-form-urlencoded" in content_type:
                            form_data = flow.request.content.decode('utf-8')
                            f.write(f"  类型: 表单数据\n")
                            f.write(f"  长度: {len(flow.request.content)} bytes\n")
                            f.write(f"  内容: {form_data}\n")
                        else:
                            f.write(f"  类型: {content_type or '未知'}\n")
                            f.write(f"  长度: {len(flow.request.content)} bytes\n")
                            # 尝试显示文本内容
                            try:
                                text_content = flow.request.content.decode('utf-8')
                                if len(text_content) > 500:
                                    f.write(f"  内容: {text_content[:500]}...[截断]\n")
                                else:
                                    f.write(f"  内容: {text_content}\n")
                            except:
                                f.write(f"  内容: [二进制数据，已跳过]\n")
                    except Exception as e:
                        f.write(f"  解析失败: {str(e)}\n")
                        f.write(f"  原始长度: {len(flow.request.content)} bytes\n")
                    f.write("\n")
                
                # curl命令
                curl_command = self.generate_curl(flow)
                f.write("🔄 等价curl命令:\n")
                f.write(f"{curl_command}\n")
                
                f.write("\n" + "=" * 80 + "\n\n")
                
        except Exception as e:
            print(f"❌ 保存请求失败: {e}")
            ctx.log.error(f"保存请求失败: {e}")
    
    def generate_curl(self, flow: mitmproxy.http.HTTPFlow) -> str:
        """生成curl命令"""
        method = flow.request.method
        url = flow.request.pretty_url
        
        # 构建curl命令
        curl_parts = ["curl"]
        
        # 添加请求方法
        if method != "GET":
            curl_parts.extend(["-X", method])
        
        # 添加请求头
        for name, value in flow.request.headers.items():
            # 跳过一些自动生成的头部
            if name.lower() in ['host', 'content-length']:
                continue
            curl_parts.extend(["-H", f"'{name}: {value}'"])
        
        # 添加请求体
        if flow.request.content:
            content_type = flow.request.headers.get("content-type", "").lower()
            try:
                if "application/json" in content_type:
                    # JSON数据
                    json_data = json.loads(flow.request.content.decode('utf-8'))
                    json_str = json.dumps(json_data, ensure_ascii=False, separators=(',', ':'))
                    curl_parts.extend(["-d", f"'{json_str}'"])
                elif "application/x-www-form-urlencoded" in content_type:
                    # 表单数据
                    form_data = flow.request.content.decode('utf-8')
                    curl_parts.extend(["-d", f"'{form_data}'"])
                else:
                    # 其他类型的数据
                    try:
                        data = flow.request.content.decode('utf-8')
                        if len(data) > 1000:
                            data = data[:1000] + "...[truncated]"
                        curl_parts.extend(["-d", f"'{data}'"])
                    except:
                        curl_parts.append("# [包含二进制数据，请查看上方请求信息]")
            except Exception:
                curl_parts.append("# [请求体解析失败]")
        
        # 添加其他有用选项
        curl_parts.extend([
            "-k",  # 忽略SSL证书错误
            "--compressed",  # 支持压缩
            "--location",  # 跟随重定向
            "-v"  # 详细输出
        ])
        
        # 最后添加URL
        curl_parts.append(f"'{url}'")
        
        return " \\\n  ".join(curl_parts)
    
    def response(self, flow: mitmproxy.http.HTTPFlow) -> None:
        """处理HTTP响应（可选，用于添加响应状态）"""
        try:
            with open(self.output_file, "a", encoding="utf-8") as f:
                if flow.response:
                    f.write(f"📥 响应状态: {flow.response.status_code} {flow.response.reason}\n")
                    f.write(f"   响应时间: {datetime.now().strftime('%H:%M:%S')}\n")
                else:
                    f.write(f"📥 响应状态: 无响应\n")
                f.write(f"   URL: {flow.request.pretty_url}\n\n")
        except Exception as e:
            ctx.log.error(f"记录响应状态失败: {e}")
    
    def done(self):
        """脚本结束时调用"""
        message = f"✅ 导出完成！共记录了 {self.request_count} 个HTTP请求"
        print(message)
        print(f"📁 文件位置: {os.path.abspath(self.output_file)}")
        ctx.log.info(message)
        ctx.log.info(f"文件保存到: {os.path.abspath(self.output_file)}")

# 创建导出器实例
addons = [SimpleCurlExporter()]